/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package cn.ac.iie.ulss.data625.handler.datadispatch;

import cn.ac.iie.ulss.data625.commons.RuntimeEnv;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;
import org.apache.log4j.Logger;
import org.apache.log4j.PropertyConfigurator;

/**
 *
 * @author Carlisle
 */
public class TokenBucketManager {

    private static Map<String, AtomicLong> tokenBucket = new ConcurrentHashMap<String, AtomicLong>();
    private static Map<String, AtomicLong> updateRate = new ConcurrentHashMap<String, AtomicLong>();
    private static Map<String, Long> globalUpdateRate = new ConcurrentHashMap<String, Long>();
    private static Map<String, Long> bucketSize = new ConcurrentHashMap<String, Long>();

    private static ZkClient zc = null;
    private static Logger logger = null;

    static {
        PropertyConfigurator.configure("log4j.properties");
        logger = Logger.getLogger(TokenBucketManager.class.getName());
    }

    public TokenBucketManager() {

    }

    public static boolean init() {
        try {
            if (!RuntimeEnv.zc.exists("/ulss/dd/tokenbucket")) {
                RuntimeEnv.zc.createPersistent("/ulss/dd/tokenbucket");
            }
            List<String> schemaList = RuntimeEnv.zc.getChildren("/ulss/dd/tokenbucket");

            if (schemaList.isEmpty()) {
                logger.info("No flow conctrol schema found in /ulss/dd/tokenbucket");
            } else {
                for (String s : schemaList) {
                    Long rate = getValueFromZK("/ulss/dd/tokenbucket/" + s + "/updaterate");
                    Long size = getValueFromZK("/ulss/dd/tokenbucket/" + s + "/bucketsize");

                    updateRate.put(s, new AtomicLong(rate / RuntimeEnv.getLoadClusterSize()));
                    tokenBucket.put(s, new AtomicLong(size));
                    globalUpdateRate.put(s, rate);
                    bucketSize.put(s, size);

                    logger.debug("update rate for " + s + " is " + rate / RuntimeEnv.getLoadClusterSize());

                    RuntimeEnv.zc.subscribeDataChanges("/ulss/dd/tokenbucket/" + s + "/updaterate", new TokenBucketWatcher());
                    RuntimeEnv.zc.subscribeDataChanges("/ulss/dd/tokenbucket/" + s + "/bucketsize", new TokenBucketWatcher());
                }

                Timer updater = new Timer("token-updater");
                Long updatePeriod = getValueFromZK("/ulss/dd/tokenbucket");
                if (updatePeriod == null || updatePeriod <= 0) {
                    logger.error("【" + updatePeriod + "】 is not a valid value, please check data of zookeeper path \"/ulss/dd/tokenbucket\" .");
                    System.exit(0);
                }

                updater.scheduleAtFixedRate(new TokenBucketUpdater(), 0, updatePeriod);
            }
            logger.info(TokenBucketManager.class.getName() + " initialized successfully.");
            return true;
        } catch (Exception ex) {
            ex.printStackTrace();
            logger.error("Initialize tokenbucket manager failed for " + ex.getMessage());
            return false;
        }
    }

    public static Integer getTokenNum(String schema) {
        return tokenBucket.get(schema).intValue();
    }

    //method given to dispatcher to consume tokens, 
    //if schema is not in tokenbucket, it doesn't need to flow control, cause schema is verified to be known.
    //otherwise, check if token is enough.
    public static boolean consumeTokenIfEnough(String schema, int datasize) {
        if (!tokenBucket.containsKey(schema)) {
            return true;
        }
        if (tokenBucket.get(schema).intValue() < datasize) {
            logger.debug("schema: " + schema + " token: " + tokenBucket.get(schema).intValue() + " is less than data size " + datasize + ", flow control.");
            return false;
        } else {
            tokenBucket.get(schema).addAndGet((-1) * datasize);
            logger.debug("schema: " + schema + " token: " + tokenBucket.get(schema).intValue() + " after consume.");
            return true;
        }
    }

    //method given to update update rate of each schema when cluster node number changes
    public static void reCalculateUpdateRate(int nodeNum) {
        List<String> schemaList = RuntimeEnv.zc.getChildren("/ulss/dd/tokenbucket");
        for (String s : schemaList) {
            Long rate = globalUpdateRate.get(s) / RuntimeEnv.getLoadClusterSize();
            updateRate.put(s, new AtomicLong(rate));
            logger.debug("Node update rate for [" + s + "] changed to " + rate);
        }
    }

    public static Long getValueFromZK(String path) {
        Object o = RuntimeEnv.zc.readData(path);
        if ((o instanceof String) || (o instanceof Long) || (o instanceof Integer)) {
            o = (o instanceof String) ? Long.decode((String) o) : o;
            o = (o instanceof Integer) ? ((Integer) o).longValue() : o;
            return (o == null || (Long) o < 0) ? 0 : (Long) o;
        } else {
            logger.error("Format of data in " + path + " is wrong, it must be one of String/Long/Integer.");
            return null;
        }
    }

    //TokenBucketUpdater updates token number in buckets
    private static class TokenBucketUpdater extends TimerTask {

        @Override
        public void run() {
            if (RuntimeEnv.loadClusterSizeChanged) {
                logger.debug("Loadcluster size changed, will recalculate update rate.");
                reCalculateUpdateRate(RuntimeEnv.getLoadClusterSize());
                RuntimeEnv.loadClusterSizeChanged = false;
            }
            Long inc = null;
            //Long starttime = System.currentTimeMillis();
            for (String schema : updateRate.keySet()) {
                inc = updateRate.get(schema).longValue();
                //check whether bucket will be full after update this time, if so, set bucket number to max size
                if (tokenBucket.get(schema).longValue() + inc > bucketSize.get(schema).longValue()) {
                    tokenBucket.get(schema).set(bucketSize.get(schema).longValue());
                    //logger.debug("bucket of " + schema + " now size is " + tokenBucket.get(schema));
                } else {
                    tokenBucket.get(schema).addAndGet(inc);
                    //logger.debug("bucket of " + schema + " now size is " + tokenBucket.get(schema));
                }
            }
            //logger.debug(updateRate.size() + " schemas updated: " + (System.currentTimeMillis() - starttime) + "ms taken.");
        }
    }

    //Watcher of token bucket, update value of specific variable when data of znode changes.
    private static class TokenBucketWatcher implements IZkDataListener {

        @Override
        public void handleDataChange(String path, Object o) throws Exception {
            String schema = Arrays.asList(path.split("/")).get(4);
            if (path.endsWith("updaterate")) {
//                updateRate.get(schema).set(Long.decode((String) o));
                updateRate.get(schema).set((Long) o);
                logger.info("Gloabal update rate for [" + schema + "] changed to " + o);
            } else {
//                if (Long.decode((String) o) < tokenBucket.get(schema).longValue()) {
                if ((Long) o < tokenBucket.get(schema).longValue()) {
                    logger.warn("Trying to modify bucket size of [" + schema + "] ,but current value is less than token number, system will not change bucket size and previous size will be write back to znode.");
//                    RuntimeEnv.zc.writeData(path, String.valueOf(bucketSize.get(schema)));
                    RuntimeEnv.zc.writeData(path, bucketSize.get(schema));
                } else {
                    bucketSize.put(schema, (Long) o);
                    logger.info("Bucket size for [" + schema + "] changed to " + o);
                }
            }
        }

        @Override
        public void handleDataDeleted(String string) throws Exception {
            logger.warn("data deleted: " + string);
        }
    }
}
